#include "PbrProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "Light.h"

void PbrProcess::OnInit()
{
	mExposure = 2.5f;
	mGamma = 2.2f;

	mWatchMode = 0;

	mLightEntitys = mGameScene->GetEntitysByType<LocalEntity, Light>();

	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();
	ubo.model = glm::scale(glm::mat4(1.0), glm::vec3(6.f));
	ubo.camera = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.f);

	LightUbo lightUbo;
	for (uint32 i = 0; i < mLightEntitys.size(); ++i)
	{
		Light& light = mLightEntitys[i]->GetComponent<Light>();
		lightUbo.lights[i] = glm::vec4(mLightEntitys[i]->GetTransform().GetLocalPosition(), 1.f);
		lightUbo.colors[i] = glm::vec4(light.mColor, light.mActive ? 1.f : 0.f);
	}
	lightUbo.exposure = mExposure;
	lightUbo.gamma = mGamma;
	lightUbo.mWatchMode = mWatchMode;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo, lightUbo](VkCommandBuffer commandBuffer)
		{
			Ubo entityUbo = ubo;

			Model model("Cerberus.obj");

			mVertexBuffer = VertexBuffer::Builder()
				.SetData(model.mVertexs.data())
				.SetSize(model.mVertexs.size() * sizeof(Vertex))
				.Build();

			mIndexBuffer = IndexBuffer::Builder()
				.SetData(model.mIndices.data())
				.SetSize(model.mIndices.size() * sizeof(uint32))
				.SetIndexCount(model.mIndices.size())
				.Build();

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&entityUbo)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			LightUbo factLightUbo = lightUbo;

			mLightBuffer = CustomBuffer::Builder()
				.SetData(&factLightUbo)
				.SetSize(sizeof(LightUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			RefCountPtr<RHITexture> irradianceCube = GetTexture("IrradianceProcess", "IrradianceCube");
			RefCountPtr<RHISampler> irradianceSampler = GetSampler("IrradianceProcess", "IrradianceSampler");

			RefCountPtr<RHITexture> brbflutTexture = GetTexture("BrbflutProcess", "ColorRenderTarget");
			RefCountPtr<RHISampler> brbflutSampler = RHISampler::Builder().Build();

			RefCountPtr<RHITexture> prefilterenvmapCube = GetTexture("PrefilterenvmapProcess", "PrefilterenvmapCube");
			RefCountPtr<RHISampler> prefilterenvmapSampler = GetSampler("PrefilterenvmapProcess", "PrefilterenvmapSampler");


			mAlbedoTexture = Picture2DTexture::Builder()
				.SetAutoMip()
				.SetPath("Pbr\\Cerberus_Albedo.jpg")
				.Build();

			mNormalTexture = Picture2DTexture::Builder()
				.SetAutoMip()
				.SetPath("Pbr\\Cerberus_Normal.jpg")
				.Build();

			mMetallicTexture = Picture2DTexture::Builder()
				.SetAutoMip()
				.SetPath("Pbr\\Cerberus_Metallic.jpg")
				.Build();

			mRoughnessTexture = Picture2DTexture::Builder()
				.SetAutoMip()
				.SetPath("Pbr\\Cerberus_Roughness.jpg")
				.Build();

			mPbrSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetMaxLod(static_cast<float>(mAlbedoTexture->GetMipLevels()))
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build();

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<PbrTextureVertexShader>()
				.PushShader<PbrTextureFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptor->SetValue("lightUbo", mLightBuffer);
			mDescriptor->SetValue("samplerIrradiance", irradianceCube, irradianceSampler);
			mDescriptor->SetValue("samplerBRDFLUT", brbflutTexture, brbflutSampler);
			mDescriptor->SetValue("prefilteredMap", prefilterenvmapCube, prefilterenvmapSampler);

			mDescriptor->SetValue("albedoMap", mAlbedoTexture, mPbrSampler);
			mDescriptor->SetValue("normalMap", mNormalTexture, mPbrSampler);
			mDescriptor->SetValue("metallicMap", mMetallicTexture, mPbrSampler);
			mDescriptor->SetValue("roughnessMap", mRoughnessTexture, mPbrSampler);

			mDescriptorMask = mDescriptor->AllocateDescriptor();

			RefCountPtr<RHITexture> colorRenderTarget = GetTexture("LocalProcess", "ColorRenderTarget");
			RefCountPtr<RHITexture> depthRenderTarget = GetTexture("LocalProcess", "DepthRenderTarget");
			RefCountPtr<RHIRenderPass> renderPass = GetRenderPass("LocalProcess", "RenderPass");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(colorRenderTarget)
				.SetDepthRenderTarget(depthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetRenderPass(renderPass->GetRenderPass())
				.Build();
		});
}

void PbrProcess::OnTick(float deltaTime)
{
	Ubo ubo;
	ubo.proj = mGameScene->GetCamera().GetProj();
	ubo.view = mGameScene->GetCamera().GetView();
	ubo.model = glm::scale(glm::mat4(1.0), glm::vec3(6.f));
	ubo.camera = glm::vec4(mGameScene->GetCamera().GetPosition(), 1.f);

	LightUbo lightUbo;
	for (uint32 i = 0; i < mLightEntitys.size(); ++i)
	{
		Light& light = mLightEntitys[i]->GetComponent<Light>();
		lightUbo.lights[i] = glm::vec4(mLightEntitys[i]->GetTransform().GetLocalPosition(), 1.f);
		lightUbo.colors[i] = glm::vec4(light.mColor, light.mActive ? 1.f : 0.f);
	}
	lightUbo.exposure = mExposure;
	lightUbo.gamma = mGamma;
	lightUbo.mWatchMode = mWatchMode;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo, lightUbo](VkCommandBuffer commandBuffer)
		{
			Ubo entityUbo = ubo;

			mUboBuffer->WriteMemory(&entityUbo);

			LightUbo factLightUbo = lightUbo;

			mLightBuffer->WriteMemory(&factLightUbo);

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkBuffer vertexBuffer = mVertexBuffer->GetBuffer();
			VkDeviceSize offsets = 0;
			vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

			VkBuffer indexBuffer = mIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

			vkCmdDrawIndexed(commandBuffer, mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void PbrProcess::OnResize()
{
}
